# 前12+12bit为水印大小信息
from tkinter import *
import tkinter.filedialog
from PIL import Image


# bin()函数将十进制转化为二进制，但会在头部加‘0b’，使用replace用空字符替换‘0b’，缺点为自动忽略高位的0，使用plus补位
def plus(_str, n):
    # print(len(_str))
    return _str.zfill(n + 2)


# 处理水印信息（图片）
# 返回各个像素点的二进制rgb组成的集合
def getcode(watermark):
    _str = ''
    # 水印大小
    _str = _str + plus(bin(watermark.size[0]), 12).replace('0b', '') + plus(
        bin(watermark.size[1]), 12).replace('0b', '')
    for i in range(watermark.size[0]):
        for j in range(watermark.size[1]):
            rgb = watermark.getpixel((i, j))
            # print(rgb)
            _str = _str + plus(bin(rgb[0]), 8).replace('0b', '')
            _str = _str + plus(bin(rgb[1]), 8).replace('0b', '')
            _str = _str + plus(bin(rgb[2]), 8).replace('0b', '')
    # print(_str)
    return _str


# 处理原图片
def LSB(img, code):
    # 计数
    count = 0
    # 水印信息的 长度
    codelen = len(code) - 1
    # print(codelen)
    for i in range(img.size[0]):
        for j in range(img.size[1]):
            data = img.getpixel((i, j))
            # 原图片每个像素的rgb信息
            r = data[0]
            g = data[1]
            b = data[2]
            # 添加水印消息
            r = (r - r % 2) + int(code[count])
            if count == codelen:
                img.putpixel((i, j), (r, g, b))
                break
            count += 1
            g = (g - g % 2) + int(code[count])
            if count == codelen:
                img.putpixel((i, j), (r, g, b))
                break
            count += 1
            b = (b - b % 2) + int(code[count])
            if count == codelen:
                img.putpixel((i, j), (r, g, b))
                break
            count += 1
            img.putpixel((i, j), (r, g, b))
            # print(count)
    fname = tkinter.filedialog.asksaveasfilename(title=u'保存文件',
                                                 filetypes=[("PNG", ".png")])
    img.save(fname + '.png', 'PNG')


def DLSB(img):
    # 水印长度
    length = ''
    # 水印宽度
    height = ''
    # 计数
    count = 0
    # 水印信息
    wt = ''
    for i in range(img.size[0]):
        for j in range(img.size[1]):
            data = img.getpixel((i, j))
            # print(count)
            # 水印长度和宽度
            if (count < 4):
                length = length + str(data[0] % 2) + str(data[1] % 2) + str(
                    data[2] % 2)
                count += 1
                continue
            elif (count < 8):
                height = height + str(data[0] % 2) + str(data[1] % 2) + str(
                    data[2] % 2)
                count += 1
                if (count == 8):
                    length = int(length, 2)
                    height = int(height, 2)
                    # print(length)
                    # print(height)
                continue
            # 提取水印信息
            # 全部水印提取完毕
            if (count == (length * height) * 8 + 12 + 12):
                break
            wt = wt + str(data[0] % 2) + str(data[1] % 2) + str(data[2] % 2)
            count += 1
    Show(length, height, wt)


def Show(length, height, wt):
    _str = []
    for i in range(0, len(wt), 8):
        # 以每8位为一组二进制，转换为十进制
        _str.append(int(wt[i:i + 8], 2))
    # 绘制和显示水印图片
    img_out = Image.new("RGB", (length, height))
    flag = 0
    for i in range(0, length):
        for j in range(0, height):
            img_out.putpixel((i, j), (_str[flag], _str[flag + 1], _str[flag + 2]))
            flag += 3
    fname = tkinter.filedialog.asksaveasfilename(title=u'保存文件',
                                                 filetypes=[("PNG", ".png")])
    img_out.save(fname + '.png', 'PNG')


# 选择图片
def choose():
    file_path = tkinter.filedialog.askopenfilename()
    return file_path


root = Tk()  # 创建一个Tkinter.Tk()实例
root.withdraw()  # 将Tkinter.Tk()实例隐藏

print("【1】 添加水印")
print("【2】 提取水印")
select = input("请选择：")
if (select == '1'):
    im = Image.open(choose())
    watermark = Image.open(choose())
    code = getcode(watermark)
    LSB(im, code)
elif (select == '2'):
    img = Image.open(choose())
    DLSB(img)
else:
    print("选择错误")
